#include <iostream>
#include <poll.h>
#include <sys/timeb.h>
#include <cstdlib>

#include "tcetools/socket.h"

using namespace std;
using namespace tcetools;

struct timeb start;

inline void stopwatch_start() {
	ftime(&start);
}

int stopwatch_stop() {
	struct timeb stop;
	ftime(&stop);
	
	int sec = stop.time - start.time;
	return (sec * 1000) + (((int) stop.millitm) - ((int) start.millitm));
}

void run_client(const char *address, const char *port) {
	cout << "Running client side latency test\n" 
		 << "Server: " << address << ':' << port << endl;
	ClientSocket cs(address, port);
	cout << "Connected, sending data." << endl;
	
	char buf[4];
	int &ibuf = *((int *) buf);
	unsigned int time_cul = 0;

	struct pollfd pfd;
	pfd.fd = cs.get_fd();
	pfd.events = POLLIN;
	pfd.revents = 0;

	for (int i = 0; i < 1024; i++) {
		ibuf = rand();

		cs.blocking_write(buf, 4);
		stopwatch_start();

		poll(&pfd, 1, -1);
		if ((pfd.revents & POLLERR) ||
			(pfd.revents & POLLHUP) ||
			(pfd.revents & POLLNVAL)) {
			cout << "Listener terminated?!?" << endl;
			return;
		}

		time_cul += stopwatch_stop();

		char buf2[4];
		int &ibuf2 = *((int *) buf);
		ssize_t readlen = cs.raw_read(buf2, 4);
		if (ibuf2 != ibuf) cout << "Data returned did not match!" << endl;

	}

	cout << "Average RTT: " << (time_cul / 1024) << "ms" << endl;
}

void run_server() {
	cout << "Running server side latency test" << endl;
	ListenerSocket ls("5001", 1);
	ServerSocket &ss = *(ls.accept_connection());
	cout << "Connection received. Begining echo loop." << endl;

	struct pollfd pfd;
	pfd.fd = ss.get_fd();
	pfd.events = POLLIN;
	pfd.revents = 0;

	while (1) {
		int pret = poll(&pfd, 1, -1);
		if ((pfd.revents & POLLERR) ||
			(pfd.revents & POLLHUP) ||
			(pfd.revents & POLLNVAL)) {
			cout << "Listener terminated?!?" << endl;
			return;
		}

		char buf[4];

		if (pfd.revents & POLLIN) {
			ssize_t readlen = ss.raw_read(buf, 4);
			if (readlen == 0) {
				cout << "Connection closed" << endl;
				return;
			} else if (readlen < 0) {
				cout << "Error!" << endl;
				return;
			}

			ss.blocking_write(buf, readlen);
		}
	}
}


int main(int argc, char *argv[]) {
	if (argc < 2 || argv[1][0] != '-') {
		cout << "No command arguements given" << endl;
		return -1;
	}

	char mode = argv[1][1];
	if (mode == 's') run_server();
	else if (mode == 'c' && argc == 4) run_client(argv[2], argv[3]);
	else {
		cout << "Invalid command" << endl;
		return -1;
	}

	return 0;
}
